// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/dwio/common/file_sink.h>
#include <pollux/dwio/common/options.h>
#include <pollux/dwio/common/statistics.h>
#include <pollux/testing/dwio/data_set_builder.h>
#include <pollux/dwio/parquet/register_parquet_reader.h>
#include <pollux/dwio/parquet/reader/parquet_reader.h>
#include <pollux/dwio/parquet/writer/writer.h>
#include <pollux/common/file/temp_directory_path.h>

#include <melon/benchmark.h>
#include <melon/init/init.h>

namespace kumo::pollux::parquet::test {

constexpr uint32_t kNumRowsPerBatch = 60000;
constexpr uint32_t kNumBatches = 50;
constexpr uint32_t kNumRowsPerRowGroup = 10000;
constexpr double kFilterErrorMargin = 0.2;

class ParquetReaderBenchmark {
 public:
  explicit ParquetReaderBenchmark(
      bool disableDictionary,
      const kumo::pollux::RowTypePtr& rowType)
      : disableDictionary_(disableDictionary) {
    rootPool_ = kumo::pollux::memory::memoryManager()->addRootPool(
        "ParquetReaderBenchmark");
    leafPool_ = rootPool_->addLeafChild("ParquetReaderBenchmark");
    dataSetBuilder_ =
        std::make_unique<kumo::pollux::test::DataSetBuilder>(*leafPool_, 0);
    auto path = fileFolder_->getPath() + "/" + fileName_;
    auto localWriteFile =
        std::make_unique<kumo::pollux::LocalWriteFile>(path, true, false);
    auto sink = std::make_unique<kumo::pollux::dwio::common::WriteFileSink>(
        std::move(localWriteFile), path);
    kumo::pollux::parquet::WriterOptions options;
    if (disableDictionary_) {
      // The parquet file is in plain encoding format.
      options.enableDictionary = false;
    }
    options.memoryPool = rootPool_.get();
    writer_ = std::make_unique<kumo::pollux::parquet::Writer>(
        std::move(sink), options, rowType);
  }

  ~ParquetReaderBenchmark() {}

  void writeToFile(
      const std::vector<kumo::pollux::RowVectorPtr>& batches,
      bool /*forRowGroupSkip*/);

  kumo::pollux::dwio::common::FilterSpec createFilterSpec(
      const std::string& columnName,
      float startPct,
      float selectPct,
      const kumo::pollux::TypePtr& type,
      bool isForRowGroupSkip,
      bool allowNulls);

  std::shared_ptr<kumo::pollux::dwio::common::ScanSpec> createScanSpec(
      const std::vector<kumo::pollux::RowVectorPtr>& batches,
      kumo::pollux::RowTypePtr& rowType,
      const std::vector<kumo::pollux::dwio::common::FilterSpec>& filterSpecs,
      std::vector<uint64_t>& hitRows);

  std::unique_ptr<kumo::pollux::dwio::common::RowReader> createReader(
      std::shared_ptr<kumo::pollux::dwio::common::ScanSpec> scanSpec,
      const kumo::pollux::RowTypePtr& rowType);

  // This method is the place where we do the read opeartions.
  // scanSpec contains the setting of filters. e.g.
  // filterRateX100 = 30 means it would filter out 70% of rows and 30% remain.
  // nullsRateX100 = 70 means it would filter out 70% of rows and 30% remain.
  // Return the number of rows after the filter and null-filter.
  int read(
      const kumo::pollux::RowTypePtr& rowType,
      std::shared_ptr<kumo::pollux::dwio::common::ScanSpec> scanSpec,
      uint32_t nextSize);

  void readSingleColumn(
      const std::string& columnName,
      const kumo::pollux::TypePtr& type,
      float startPct,
      float selectPct,
      uint8_t nullsRateX100,
      uint32_t nextSize);

 private:
  const std::string fileName_ = "test.parquet";
  const std::shared_ptr<kumo::pollux::exec::test::TempDirectoryPath>
      fileFolder_ = kumo::pollux::exec::test::TempDirectoryPath::create();
  const bool disableDictionary_;

  std::unique_ptr<kumo::pollux::test::DataSetBuilder> dataSetBuilder_;
  std::shared_ptr<kumo::pollux::memory::MemoryPool> rootPool_;
  std::shared_ptr<kumo::pollux::memory::MemoryPool> leafPool_;
  std::unique_ptr<kumo::pollux::parquet::Writer> writer_;
  kumo::pollux::dwio::common::RuntimeStatistics runtimeStats_;
};

void run(
    uint32_t,
    const std::string& columnName,
    const kumo::pollux::TypePtr& type,
    float filterRateX100,
    uint8_t nullsRateX100,
    uint32_t nextSize,
    bool disableDictionary);

} // namespace kumo::pollux::parquet::test
